package quantum.state;
import org.apache.commons.math.complex.Complex;
import org.apache.commons.math.exception.OutOfRangeException;

import quantum.sampling.RejectionSampler;

public class QuantumState {
	Complex[] m_amplitudes;
	
	public QuantumState(Complex[] amplitudes) {
		m_amplitudes = amplitudes;
	}
	
	public QuantumState(int numQubits) {
		m_amplitudes = new Complex[1 << numQubits];
	}
	
	public Complex getAmplitude(int i) throws OutOfRangeException {
		if( i >= m_amplitudes.length )
			throw new OutOfRangeException(i, 0, m_amplitudes.length - 1);
		
		return m_amplitudes[i];
	}
	
	public void setAmplitude(int i, Complex c) throws OutOfRangeException {
		if( i >= m_amplitudes.length )
			throw new OutOfRangeException(i, 0, m_amplitudes.length - 1);
		
		m_amplitudes[i] = c; 
	}
	
	public int measure() {
		return RejectionSampler.sample(generate_probabilities());
	}
	
	public Complex[] toVector() {
		return m_amplitudes.clone();
	}
	
	public void normalize() {
		
	}
	
	public int numAmplitudes() {
		return m_amplitudes.length;
	}
	
	private double[] generate_probabilities() {
		double[] probs = new double[m_amplitudes.length];
		for(int i = 0 ; i < probs.length ; i++) {
			double norm = m_amplitudes[i].abs();
			probs[i] = norm*norm;
		}
		
		return probs;
	}
}
